#!/bin/bash
# DDPM训练脚本 - 使用: bash train_ddpm.sh
export CUDA_VISIBLE_DEVICES=0

# 路径配置
DATA_DIR="/home/hlx/lsq/Squidiff/data/imputed_npz"
LOG_DIR="/home/hlx/lsq/Squidiff/logs/uncond_logs2"
CHECKPOINT_DIR="/home/hlx/lsq/Squidiff/model/uncond_checkpoints2"
# mkdir -p "${LOG_DIR}/checkpoints" "${LOG_DIR}/visualizations"


# 启动前检查
echo "=========================================="
echo "DDPM Hi-C 训练启动"
echo "=========================================="
echo "数据目录: ${DATA_DIR}"
echo "日志目录: ${LOG_DIR}"
echo "检查点目录: ${CHECKPOINT_DIR}"

echo ""
echo "GPU 状态:"
nvidia-smi --query-gpu=index,name,memory.total,memory.free --format=csv,noheader 2>/dev/null || echo "无法检测GPU"
echo ""
echo "开始训练..."
echo "=========================================="
echo ""

python train_hic_diffusion_enhanced.py \
    --data_dir "${DATA_DIR}" \
    --logger_path "${LOG_DIR}" \
    --resume_checkpoint "${CHECKPOINT_DIR}" \
    \
    --timepoints "G1,late_S1" \
    --file_format "npz" \
    --chromosome "chr1" \
    --resolution 50000 \
    --window_size 128 \
    --normalization "oe" \
    --log_transform True \
    --interpolation_prob 0.0 \
    --augmentation False \
    --use_grid_windows True \
    --band_limit_bp 5000000 \
    \
    --diffusion_steps 1000 \
    --noise_schedule "linear" \
    --learn_sigma False \
    --predict_xstart False \
    --use_kl False \
    --rescale_timesteps False \
    \
    --model_channels 128 \
    --num_res_blocks 2 \
    --attention_resolutions "32,16,8" \
    --channel_mult "1,2,4,8" \
    --dropout 0.1 \
    --use_scale_shift_norm True \
    --num_heads 4 \
    --use_fp16 False \
    \
    --time_cond_dim 128 \
    --use_transition_anchors False \
    --cond_drop_prob 0.1 \
    \
    --batch_size 4 \
    --microbatch -1 \
    --lr 1e-4 \
    --lr_anneal_steps 500000 \
    --ema_rate "0.9999" \
    --schedule_sampler "uniform" \
    --log_interval 100 \
    --save_interval 5000 \
    \
    --multi_t_training False \
    --num_t_per_batch 11 \
    --t_multi_strategy "linspace" \
    --t_stride 100 \
    \
    --validation_split 0.1 \
    --validation_interval 5000 \
    --validation_samples 4 \
    \
    2>&1 | tee "${LOG_DIR}/training.log"

echo "训练完成！日志: ${LOG_DIR}"

